'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 2_study_qa_extra_with_json.py
* @Time: 2025/10/29
* @All Rights Reserve By Brtc
'''
import dotenv
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI

from pydantic import Field, BaseModel

dotenv.load_dotenv()

class QAExtra(BaseModel):
    """一个问答对工具。传递对应的假设性问题 + 答案"""
    question:str = Field(description="假设性问题")
    answer:str = Field(description="假设性答案")

llm = ChatOpenAI(model = "gpt-4o")
struct_llm = llm.with_structured_output(QAExtra, method = "json_mode")

prompt = ChatPromptTemplate.from_messages([
    ("system", "请从用户传递的query中提取假设性的问题 + 答案，响应格式是JSON， 并携带'question 和 'answer两个字段。'"),
    ("human", "{query}")
])
chain = {"query":RunnablePassthrough()}| prompt | struct_llm
res = chain.invoke("今天星期三")
print(type(res))
print(res.answer)
print(res)